import os
import time

from file_utils import (
    display_directory_tree,
    collect_file_paths,
    separate_files_by_type,
    read_file_data
)

from data_processing_common import (
    compute_operations,
    execute_operations,
    process_files_by_date,
    process_files_by_type,
)

from text_data_processing import (
    process_text_files
)

from image_data_processing import (
    process_image_files
)

from output_filter import filter_specific_output  # 导入上下文管理器
from nexa.gguf import NexaVLMInference, NexaTextInference  # 导入模型类

def ensure_nltk_data():
    """确保NLTK数据已高效静默下载"""
    import nltk
    nltk.download('stopwords', quiet=True)
    nltk.download('punkt', quiet=True)
    nltk.download('wordnet', quiet=True)
    nltk.download('punkt_tab',quiet=True)

# 初始化模型
image_inference = None
text_inference = None

def initialize_models():
    """若未初始化则初始化模型"""
    global image_inference, text_inference
    if image_inference is None or text_inference is None:
        # 初始化模型
        model_path = "llava-v1.6-vicuna-7b:q4_0"
        model_path_text = "Qwen2.5-3B-Instruct:q4_K_M"

        # 使用filter_specific_output上下文管理器
        with filter_specific_output():
            # 初始化图像推理模型
            image_inference = NexaVLMInference(
                model_path=model_path,
                local_path=None,
                stop_words=[],
                temperature=0.3,
                max_new_tokens=3000,
                top_k=3,
                top_p=0.2,
                profiling=False,
                n_ctx=4096
            )

            # 初始化文本推理模型
            text_inference = NexaTextInference(
                model_path="Qwen2.5-3B-Instruct:q4_K_M",
                local_path="C:\\Users\\admin\\.cache\\nexa\\hub\\official\\Qwen2.5-3B-Instruct\\q4_K_M.gguf",
                stop_words=[],
                temperature=0.5,
                max_new_tokens=3000,  # 按需调整
                top_k=3,
                top_p=0.3,
                profiling=False,
                n_ctx=8192
            )
        print("**----------------------------------------------**")
        print("**       图像推理模型初始化完成                **")
        print("**       文本推理模型初始化完成                **")
        print("**----------------------------------------------**")

def simulate_directory_tree(operations, base_path):
    """根据操作建议模拟目录树结构"""
    tree = {}
    for op in operations:
        rel_path = os.path.relpath(op['destination'], base_path)
        parts = rel_path.split(os.sep)
        current_level = tree
        for part in parts:
            if part not in current_level:
                current_level[part] = {}
            current_level = current_level[part]
    return tree

def print_simulated_tree(tree, prefix=''):
    """打印模拟的目录树结构"""
    pointers = ['├── '] * (len(tree) - 1) + ['└── '] if tree else []
    for pointer, key in zip(pointers, tree):
        print(prefix + pointer + key)
        if tree[key]:  # 如果有子目录或文件
            extension = '│   ' if pointer == '├── ' else '    '
            print_simulated_tree(tree[key], prefix + extension)

def get_yes_no(prompt):
    """获取用户是/否响应"""
    while True:
        response = input(prompt).strip().lower()
        if response in ('yes', 'y', '是', '好的'):
            return True
        elif response in ('no', 'n', '否', '不'):
            return False
        elif response == '/exit':
            print("程序退出中...")
            exit()
        else:
            print("请输入'是'或'否'。输入'/exit'可退出程序。")

def get_mode_selection():
    """获取用户选择的模式"""
    while True:
        print("请选择文件整理模式：")
        print("1. 按内容分类")
        print("2. 按日期分类")
        print("3. 按类型分类")
        response = input("请输入1、2或3（输入'/exit'可退出程序）：").strip()
        if response == '/exit':
            print("程序退出中...")
            exit()
        elif response == '1':
            return 'content'
        elif response == '2':
            return 'date'
        elif response == '3':
            return 'type'
        else:
            print("无效选择，请输入1、2或3。输入'/exit'可退出程序。")

def main():
    # 确保NLTK数据已下载
    #ensure_nltk_data()

    # 初始设置为试运行模式
    dry_run = True

    # 静默模式说明
    print("-" * 50)
    print("**注意：静默模式会将所有输出记录到文本文件而非显示在终端**")
    silent_mode = get_yes_no("是否启用静默模式？(是/否)：")
    if silent_mode:
        log_file = 'operation_log.txt'
    else:
        log_file = None

    while True:
        # 路径配置
        if not silent_mode:
            print("-" * 50)

        # 获取输入输出路径
        input_path = input("请输入要整理的目录路径：").strip()
        while not os.path.exists(input_path):
            message = f"输入路径{input_path}不存在，请重新输入有效路径"
            if silent_mode:
                with open(log_file, 'a') as f:
                    f.write(message + '\n')
            else:
                print(message)
            input_path = input("请输入要整理的目录路径：").strip()

        # 确认输入路径
        message = f"成功加载输入路径：{input_path}"
        if silent_mode:
            with open(log_file, 'a') as f:
                f.write(message + '\n')
        else:
            print(message)
        if not silent_mode:
            print("-" * 50)

        # 默认输出路径为输入目录下的"organized_folder"
        output_path = input("请输入整理后文件的存储路径（直接回车将使用输入目录下的'organized_folder'）：").strip()
        if not output_path:
            output_path = os.path.join(os.path.dirname(input_path), 'organized_folder')

        # 确认输出路径
        message = f"输出路径已设置为：{output_path}"
        if silent_mode:
            with open(log_file, 'a') as f:
                f.write(message + '\n')
        else:
            print(message)
        if not silent_mode:
            print("-" * 50)

        # 开始处理文件
        start_time = time.time()
        file_paths = collect_file_paths(input_path)
        end_time = time.time()

        message = f"加载文件路径耗时：{end_time - start_time:.2f}秒"
        if silent_mode:
            with open(log_file, 'a') as f:
                f.write(message + '\n')
        else:
            print(message)
        if not silent_mode:
            print("-" * 50)
            print("整理前的目录结构：")
            display_directory_tree(input_path)

            print("*" * 50)

        # 模式选择循环
        while True:
            mode = get_mode_selection()

            if mode == 'content':
                # 内容分类模式
                if not silent_mode:
                    print("正在检查模型是否已下载，若未下载将立即下载")
                initialize_models()

                if not silent_mode:
                    print("*" * 50)
                    print("文件加载成功，处理可能需要几分钟时间")
                    print("*" * 50)

                # 准备统计链接类型
                link_type_counts = {'hardlink': 0, 'symlink': 0}

                # 按类型分离文件
                image_files, text_files = separate_files_by_type(file_paths)

                # 准备文本元组用于处理
                text_tuples = []
                for fp in text_files:
                    text_content = read_file_data(fp)
                    if text_content is None:
                        message = f"不支持或无法读取的文本格式：{fp}"
                        if silent_mode:
                            with open(log_file, 'a') as f:
                                f.write(message + '\n')
                        else:
                            print(message)
                        continue  # 跳过不支持的文件
                    text_tuples.append((fp, text_content))

                # 顺序处理文件
                data_images = process_image_files(image_files, image_inference, text_inference, silent=silent_mode, log_file=log_file)
                data_texts = process_text_files(text_tuples, text_inference, silent=silent_mode, log_file=log_file)

                # 准备复制和重命名
                renamed_files = set()
                processed_files = set()

                # 合并所有数据
                all_data = data_images + data_texts

                # 计算操作步骤
                operations = compute_operations(
                    all_data,
                    output_path,
                    renamed_files,
                    processed_files
                )

            elif mode == 'date':
                # 按日期处理文件
                operations = process_files_by_date(file_paths, output_path, dry_run=False, silent=silent_mode, log_file=log_file)
            elif mode == 'type':
                # 按类型处理文件
                operations = process_files_by_type(file_paths, output_path, dry_run=False, silent=silent_mode, log_file=log_file)
            else:
                print("选择的模式无效")
                return

            # 模拟并显示建议的目录结构
            print("-" * 50)
            message = "建议的目录结构："
            if silent_mode:
                with open(log_file, 'a') as f:
                    f.write(message + '\n')
            else:
                print(message)
                print(os.path.abspath(output_path))
                simulated_tree = simulate_directory_tree(operations, output_path)
                print_simulated_tree(simulated_tree)
                print("-" * 50)

            # 询问用户是否继续
            proceed = get_yes_no("是否继续执行这些更改？(是/否)：")
            if proceed:
                # 创建输出目录
                os.makedirs(output_path, exist_ok=True)

                # 执行实际文件操作
                message = "正在执行文件操作..."
                if silent_mode:
                    with open(log_file, 'a') as f:
                        f.write(message + '\n')
                else:
                    print(message)
                execute_operations(
                    operations,
                    dry_run=False,
                    silent=silent_mode,
                    log_file=log_file
                )

                message = "文件整理完成"
                if silent_mode:
                    with open(log_file, 'a') as f:
                        f.write("-" * 50 + '\n' + message + '\n' + "-" * 50 + '\n')
                else:
                    print("-" * 50)
                    print(message)
                    print("-" * 50)
                break  # 操作成功后退出模式选择循环
            else:
                # 询问是否尝试其他分类方式
                another_sort = get_yes_no("是否选择其他分类方式？(是/否)：")
                if another_sort:
                    continue  # 返回模式选择
                else:
                    print("操作已取消")
                    break  # 退出模式选择循环

        # 询问是否整理其他目录
        another_directory = get_yes_no("是否整理其他目录？(是/否)：")
        if not another_directory:
            break  # 退出主循环


if __name__ == '__main__':
    main()